package org.pathwaycommons.nlp;

import org.biopax.paxtools.pattern.miner.SIFEnum;
import org.cbio.causality.analysis.Graph;
import org.cbio.causality.network.PathwayCommons;
import org.cbio.causality.util.FDR;

import java.io.File;
import java.io.FileNotFoundException;
import java.util.*;

/**
 * Predicts the relation type in a given table of gene pairs, based on enrichment in predefined
 * sets derived from Pathway Commons.
 */
public class DistanceLearner
{
	protected static Map<String, Set<String>> sets;

	public static void main(String[] args) throws FileNotFoundException
	{
		Set<String> table = new HashSet<String>();
		Scanner sc = new Scanner(new File("table.txt"));
		while (sc.hasNextLine())
		{
			String[] token = sc.nextLine().split("\t");
			table.add(token[0] + " " + token[1]);
		}

		final Map<String, Double>[] vals = getEnrichedTerms(table);
		List<String> terms = new ArrayList<String>(vals[0].keySet());
		Collections.sort(terms, new Comparator<String>()
		{
			@Override
			public int compare(String o1, String o2)
			{
				return vals[0].get(o1).compareTo(vals[0].get(o2));
			}
		});

		for (String term : terms)
		{
			double pval = vals[0].get(term);
			double qval = vals[1].get(term);

			if (pval > 0.05) return;

			System.out.println(term + "\t" + pval + "\t" + qval + "\t" +
				new ArrayList<String>(getOverlap(table, sets.get(term))));
		}
	}

	/**
	 * Gets the enriched terms for the given gene pairs against the available gene pair sets.
	 * Enrichment p-vals and q-vals are returned.
	 */
	public static Map<String, Double>[] getEnrichedTerms(Set<String> test)
	{
		EnrichmentTester et = new EnrichmentTester(sets, test);
		Map<String, Double> pvals = et.calcPvals();
		return new Map[]{pvals, FDR.getQVals(pvals, null)};
	}

	private static Set<String> getOverlap(Set<String> s1, Set<String> s2)
	{
		Set<String> set = new HashSet<String>(s2);
		set.retainAll(s1);
		return set;
	}

	/**
	 * Prepares the pairs sets to test the enrichment against.
	 */
	private static void prepareSets()
	{
		sets = new HashMap<String, Set<String>>();

		DeltaFeatureReader dfr = new DeltaFeatureReader();
		Map<String, Set<String>> featmap = dfr.readFeatureSets(true, true);
		for (String feat : featmap.keySet())
		{
			sets.put("gained: " + feat, featmap.get(feat));
		}
		featmap = dfr.readFeatureSets(false, true);
		for (String feat : featmap.keySet())
		{
			sets.put("lost: " + feat, featmap.get(feat));
		}

		putPairsFromGraph(sets, SIFEnum.CONTROLS_STATE_CHANGE_OF);
		putPairsFromGraph(sets, SIFEnum.CONTROLS_PHOSPHORYLATION_OF);
		putPairsFromGraph(sets, SIFEnum.CONTROLS_TRANSPORT_OF);
		putPairsFromGraph(sets, SIFEnum.CONTROLS_EXPRESSION_OF);
		putPairsFromGraph(sets, SIFEnum.CATALYSIS_PRECEDES);
		putPairsFromGraph(sets, SIFEnum.IN_COMPLEX_WITH);
	}

	private static void putPairsFromGraph(Map<String, Set<String>> map, SIFEnum sifType)
	{
		map.put(sifType.getTag(), parsePairs(PathwayCommons.getGraph(sifType)));
	}

	private static Set<String> parsePairs(Graph graph)
	{
		Set<String> pairs = new HashSet<String>();

		for (String source : graph.getOneSideSymbols(true))
		{
			for (String target : graph.getDownstream(source))
			{
				pairs.add(source + " " + target);
			}
		}

		return pairs;
	}

	static
	{
		prepareSets();
	}
}